import json
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from urllib.parse import parse_qs

class DownloadProgressConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.task_type = self.scope['url_route']['kwargs']['task_type']
        self.room_group_name = f'download_progress_{self.task_type}'
        
        # 从URL参数获取token进行认证
        query_string = self.scope.get('query_string', b'').decode()
        query_params = parse_qs(query_string)
        token = query_params.get('token', [None])[0]
        
        print(f"WebSocket连接尝试 - task_type: {self.task_type}, token: {token[:20] if token else 'None'}...")
        
        if not token:
            print("WebSocket连接失败: 缺少token")
            await self.close(code=4001)
            return
            
        # 验证token
        user = await self.get_user_from_token(token)
        if not user:
            print(f"WebSocket连接失败: token验证失败 - {token[:20]}...")
            await self.close(code=4001)
            return
            
        print(f"WebSocket连接成功 - 用户: {user.username}, task_type: {self.task_type}")
        self.scope['user'] = user
        
        # 加入房间组
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        
        await self.accept()
    
    async def disconnect(self, close_code):
        # 离开房间组
        await self.channel_layer.group_discard(
            self.room_group_name,
            self.channel_name
        )
    
    async def receive(self, text_data):
        # 处理来自WebSocket的消息（如果需要）
        pass
    
    @database_sync_to_async
    def get_user_from_token(self, token_key):
        from rest_framework.authtoken.models import Token
        try:
            token = Token.objects.get(key=token_key)
            return token.user
        except Token.DoesNotExist:
            return None
    
    # 接收来自房间组的消息
    async def download_progress_update(self, event):
        progress_data = event['progress_data']
        
        # 发送消息到WebSocket
        await self.send(text_data=json.dumps({
            'type': 'progress_update',
            'task_type': self.task_type,
            'data': progress_data
        }))